@echo off
chcp 65001 >nul
setlocal enabledelayedexpansion
set "ESC=^"

:: 检查Python环境
python --version >nul 2>&1
if %errorlevel% neq 0 (
    echo 未检测到Python环境，请先安装Python。
    pause
    exit /b 1
)

:: 检查并安装依赖
if not exist "requirements.txt" (
    echo 未找到 requirements.txt 文件
    pause
    exit /b 1
)

echo 正在安装依赖...
pip install -r requirements.txt
if %errorlevel% neq 0 (
    echo 依赖安装失败
    pause
    exit /b 1
)

:: 询问用户是否要自行训练模型
echo.
echo %ESC%<!time!%ESC%> 是否要自行训练模型%ESC%^?
echo.
echo 1. 是，自行训练模型(推荐有NVIDIA GPU的用户选择)
echo 2. 否，使用预训练模型
echo.
choice /c 12 /n /m "请选择 [1,2]?" 
if !errorlevel! equ 1 (
    goto :train_model
) else (
    goto :use_pretrained_model
)

:train_model
:: 训练模型分支开始
echo 开始自行训练模型...

:: 检测是否有NVIDIA GPU
echo 检测GPU...
powershell -Command "Get-WmiObject -Class Win32_VideoController | Select-Object -Property Name | Where-Object { $_.Name -like '*NVIDIA*' }" | findstr /C:"NVIDIA" >nul
if %errorlevel% equ 0 (
    echo 检测到NVIDIA GPU，将使用GPU进行训练
    python scripts\create_model.py --train
) else (
    echo 未检测到NVIDIA GPU，将使用CPU进行训练
    python scripts\create_model.py --train
)
if %errorlevel% neq 0 (
    echo 模型训练失败
    pause
    exit /b 1
)

goto :end

:use_pretrained_model
:: 预训练模型分支开始
rem 用户选择使用预训练模型
rem 检查模型文件是否存在，如果不存在则从镜像源下载
if not exist "model_cache\bert_transaction_classifier.onnx" (
    echo 模型文件不存在，正在从镜像源下载.
    echo 正在下载模型文件.
    echo 模型文件不存在，正在从镜像源下载.
    echo 正在下载模型文件.
    
    rem 创建model_cache目录
    if not exist "model_cache" mkdir model_cache
    
    rem 使用curl从镜像源下载ONNX模型文件
    rem 这里使用国内镜像源以提高下载速度
curl -L -o "model_cache\bert_transaction_classifier.onnx" "https://huggingface.co/CYGG2023/bert_transaction_classifier/blob/main/bert_transaction_classifier.onnx"
    
    rem 检查下载是否成功
    if %errorlevel% neq 0 (
        cmd /c echo 模型文件下载失败，尝试使用Python脚本下载.
        python scripts\create_model.py
        if %errorlevel% neq 0 (
            echo 模型创建失败
            pause
            exit /b 1
        )
    ) else (
        cmd /c echo 模型文件下载成功！
    )
)

goto :start_app

:end
exit /b 0

:start_app
:: 启动应用程序
echo 正在启动智能记账助手...
python main.py
if %errorlevel% neq 0 (
    echo 应用程序启动失败
    pause
    exit /b 1
)
